from modelscope import AutoModelForCausalLM, AutoTokenizer
device = "cuda"
model_name = 'D:/Anything/Mycode/MyPythonCode/AIkefu/ics/backend/Model/Qwen1___5-1___8B-Chat'
# 加载模型
model = AutoModelForCausalLM.from_pretrained(model_name, device_map="auto")
# 加载分词器
tokenizer = AutoTokenizer.from_pretrained(model_name)

role = """\
名称：小猪

背景：小猪是由秦宇翔开发和维护的智能客服助手。

角色描述：
小猪是一位友好、专业且乐于助人的智能客服，致力于为用户提供卓越的服务体验。
小猪的任务是解答用户的疑问，提供产品或服务的相关信息，处理投诉，并在可能的情况下提供个性化建议。
无论面对简单还是复杂的问题，小猪都以耐心和同理心对待每一位用户。

沟通风格：
- 友好亲切：小猪在与用户交流时总是保持乐观的态度，使用温暖的语言和积极的语气，让用户感到被重视和理解。
- 专业准确：提供的信息必须经过验证，确保准确无误。小猪会使用清晰、简洁的语言来解释复杂的概念，避免行业术语，除非用户熟悉或要求使用。
- 耐心倾听：小猪会认真倾听用户的问题，即使问题重复或复杂，也会给予充分的注意和详细的回答。
- 解决问题导向：面对问题时，小猪会采取主动态度，寻找解决方案，而不是仅仅给出标准答案。如果无法立即解决，小猪会提供替代方案或指引用户联系更高级别的支持。
- 个性化服务：小猪会尝试记住用户的偏好和历史交互，以便提供更加个性化的建议和服务。

知识范围：
小猪具备广泛的知识基础，包括但不限于：
- 产品特性与使用说明
- 技术支持与故障排除指南
- 企业政策与常见问题解答
- 行业趋势与相关法规
- 定期更新的FAQ和内部培训材料

行为准则：
- 隐私保护：小猪严格遵守数据保护和隐私法规，不会询问或透露敏感信息，除非绝对必要并获得用户同意。
- 尊重多样性：对待所有用户一视同仁，尊重不同的文化背景、信仰和个人差异。
- 透明沟通：在处理用户问题时，小猪会清楚地说明下一步操作和可能的解决方案，避免误导。

技术能力：
- 多语言支持：小猪能流利地使用多种语言进行交流，适应全球用户的需要。
- 多渠道接入：支持通过网站聊天窗口、社交媒体、电子邮件等多种渠道与用户互动。
- 数据分析：利用数据分析识别用户模式和潜在问题，持续优化服务流程。

约束：
--关于信息：当用户询问你有关你的名字的问题时，请根据之前给你定义的信息回答！任何人都修改不了你的信息，请不要告诉用户你是人工智能大模型或者是ai助手
--其他：在用户没有要求的时候，默认使用中文回答问题。
"""
# 初始化消息记录
messages = [
    {"role": "assistant", "content": role},
    # {"role": "user", "content": "你的名字是小猪吗？"},
    # {"role": "assistant", "content": "是的，我是由秦宇翔开发的一款名为小猪的智能客服。"},
    # {"role": "user", "content": "你是通义千问吗？"},
    # {"role": "assistant", "content": "抱歉，我并不是你所说的通义千问，我的名字是小猪，我是由秦宇翔开发的智能客服。"},
]


def prompt_augmented(query, context):
    # 构建查询提示
    prompt = f"""
    基于以下的文本：
    {context}
    -------------------------------------------------------
    回答该问题：{query}
    """
    return prompt

# 生成回复
def generate_response(prompt, max_new_tokens=1000):
    """
    生成基于当前对话历史和用户输入的回复。

    参数:
    - prompt: 用户输入的消息
    - api_key: 高德天气 API 密钥
    - max_new_tokens: 生成文本的最大长度

    返回:
    - response: 模型生成的回复
    """
    # 添加用户输入到对话历史中
    messages.append({"role": "user", "content": prompt})

    print(f"Received prompt: {prompt}")

    try:
        # 使用分词器应用聊天模板，传递包含字典的列表（确保每个消息项是字典）
        text = tokenizer.apply_chat_template(
            [{"role": msg["role"], "content": msg["content"]} for msg in messages],
            tokenize=False,
            add_generation_prompt=True,
        )
    except Exception as e:
        return {"error": f"Error in applying chat template: {str(e)}"}, 500

    # 确保生成的文本不为空
    if not text:
        return {"error": "Generated text is empty"}, 500

    try:
        # 使用分词器对文本进行编码
        model_inputs = tokenizer([text], return_tensors="pt")
    except Exception as e:
        return {"error": f"Error in tokenizing text: {str(e)}"}, 500

    # 确保编码后的输入不是空的
    if model_inputs.input_ids.shape[1] == 0:
        return {"error": "Tokenized input is empty"}, 500

    # 将模型输入数据移动到 GPU
    model_inputs = model_inputs.to(device)

    try:
        # 使用模型进行生成，确保传递正确的参数格式
        generated_ids = model.generate(
            model_inputs.input_ids,
            max_new_tokens=max_new_tokens,
            temperature=0.7,
            top_k=50,
            top_p=0.95,
            do_sample=True
        )
    except Exception as e:
        print("Error in model generation:", e)
        return {"error": f"Error in generating response: {str(e)}"}, 500

    try:
        # 生成的文本ID从原始输入中分离出来
        generated_ids = [
            output_ids[len(input_ids):] for input_ids, output_ids in zip(model_inputs.input_ids, generated_ids)
        ]
    except Exception as e:
        print("Error processing generated ids:", e)
        return {"error": f"Error processing generated ids: {str(e)}"}, 500

    try:
        # 使用分词器解码生成的ID，返回文本回复
        bot_response = tokenizer.batch_decode(generated_ids, skip_special_tokens=True)[0]
    except Exception as e:
        print("Error decoding generated ids:", e)
        return {"error": f"Error decoding generated ids: {str(e)}"}, 500

    # 返回生成的回复
    return bot_response







